# -*- coding: utf-8 -*-
# FileName:     instruments.py
# time:         23/4/23 023 下午 1:20
# Author:       Zhou Hang
# Description:  Simulate Instruments!
import math
import random
from collections import deque
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

from src.utils.zaudio import StdAudio
from src.utils.zplot import ZPlot
from src.utils.common import *


class WaveBuffer:
    def __init__(self, init_capacity):
        """
        :param init_capacity: 一个周期的采样点数量
        """
        self.buffer = deque()
        self.init_capacity = init_capacity  # 每个周期多少数据点
        for i in range(self.init_capacity):
            self.buffer.append(0)

    def sample(self):
        return self.buffer[0]

    def dump(self, cut_length=3):
        """
        将buffer作为ndarray导出
        :return:
        """
        return np.array(list(self.buffer))[::cut_length], np.arange(self.init_capacity)[::cut_length]

    def tic(self):
        pass

    def pluck(self):
        pass

    def gen_wave_pic(self, picdir: Path, what_instrument: str, frequency, T: int = 1, ylim=(-0.5, 0.5)):
        """
        将当前波形的第一帧静态图保存到文件中
        :param picdir: 待保存的文件存储的位置
        :param what_instrument: 是什么乐器, 任意字符串即可
        :param frequency: 频率
        :param T: 绘制多少个周期
        :return:
        """
        wave_datas = np.array([])

        plt.figure(num=1, figsize=(7, 4))
        title = f'{what_instrument} f={frequency}'
        plt.ylim(*ylim)
        # plt.grid()
        plt.xlabel("Sample Time Slice")
        plt.ylabel(f"Wave Value {str(ylim)}")
        for t in range(T):
            wave_data, time_slice = self.dump()
            self.tic()
            wave_datas = np.hstack((wave_datas, wave_data))
        time_slice = np.linspace(0, T, len(wave_datas))
        plt.plot(time_slice, wave_datas, 'g')
        plt.savefig(str(picdir / f'{title}.png'))

    def gen_wave_ani_super(self, anidir: Path, what_instrument: str, frequency, length=80000, z_sample_length=1000):
        """
        将当前波形的动态图保存到文件中
        :param anidir: 待保存的动画文件存储的位置
        :param what_instrument: 是什么乐器, 任意字符串即可
        :param frequency: 频率
        :param length: 发生时间长度，不同的频率该值实际对应的时间长度也不同
        :param z_sample_length: 每隔多长时间采样一次来绘图
        :return:
        """
        wave_datas = []
        wave_data, time_slice = self.dump()
        wave_datas.append(wave_data)
        fig = plt.figure(num=1, figsize=(7, 5))  # 创建图
        title = f'{what_instrument} f={frequency}'

        for i in range(length):
            self.tic()
            if i % z_sample_length == 0:
                wave_data, time_slice = self.dump()
                wave_datas.append(wave_data)

        def update(n):  # 更新函数
            plt.cla()
            plt.ylim(-0.5, 0.5)
            plt.grid()
            plt.xlabel("Sample Time Slice")
            plt.ylabel("Wave Value [-0.5 -- 0.5]")
            plt.title(f'{title} WaveForm')
            cur_wave_data = wave_datas[n]
            plt.plot(time_slice, cur_wave_data, 'g')

        ani = FuncAnimation(fig, update,
                            frames=np.arange(len(wave_datas)),
                            interval=50,
                            blit=False,
                            repeat=False)
        ani.save(str(anidir / f'{title}.mp4'), writer='ffmpeg')


class CommonWave(WaveBuffer):
    def __init__(self, init_capacity, loudness=0.5, decay=0.996):
        """
        :param init_capacity: 一个周期的采样点数量
        :param loudness: 响度
        :param decay: 声音是否自然消散
        """
        super(CommonWave, self).__init__(init_capacity)
        self.loudness = loudness
        self.decay = decay

    def tic(self):
        a = self.buffer.popleft()
        a *= self.decay
        self.buffer.append(a)


class SquareWave(CommonWave):
    def __init__(self, dutyratio=0.5, frequency=440.0, loudness=0.3, samplerate=SAMPLE_RATE, decay=0.996):
        """
        :param dutyratio: 占空比
        :param frequency: 1秒内有多少周期
        :param samplerate: 1秒内有多少数据点
        :param loudness: 响度
        :param decay: 声音是否自然消散
        """
        super(SquareWave, self).__init__(round(samplerate / frequency), loudness, decay)
        self.duty_ratio = dutyratio
        self.frequency = frequency

    def pluck(self):
        self.buffer.clear()
        # 生成方波
        for i in range(self.init_capacity):
            self.buffer.append(self.loudness if i < self.init_capacity * self.duty_ratio else 0)

    def gen_wave_ani(self, anidir: str):
        super(SquareWave, self).gen_wave_ani_super(Path(anidir),
                                                   "SquareWave",
                                                   self.frequency,
                                                   80000, 1000)


class TriangularWave(CommonWave):
    def __init__(self, topratio=0.5, frequency=440.0, loudness=0.5, samplerate=SAMPLE_RATE, decay=0.996):
        """
        :param topratio: 最高点在一个周期内到达的时间，1表示在周期结束时到达最高点
        :param frequency: 1秒内有多少周期
        :param samplerate: 1秒内有多少数据点
        :param loudness: 响度
        :param decay: 声音是否自然消散
        """
        super(TriangularWave, self).__init__(round(samplerate / frequency), loudness, decay)
        self.upnum = round(topratio * self.init_capacity)  # 上升沿数据点数量
        self.downnum = self.init_capacity - self.upnum  # 下降沿数据点数量
        self.updelta = 0 if self.upnum == 0 else self.loudness / self.upnum  # 上升沿每个数据点间隔
        self.downdelta = 0 if self.downnum == 0 else self.loudness / self.downnum  # 下降沿每个数据点间隔
        self.frequency = frequency

    def pluck(self):
        # 生成三角波
        self.buffer.clear()
        for i in range(self.upnum):
            self.buffer.append(self.updelta * i)
        for i in range(self.upnum, self.init_capacity):
            self.buffer.append(self.loudness - self.downdelta * (i - self.upnum))

    def gen_wave_ani(self, anidir: str):
        super(TriangularWave, self).gen_wave_ani_super(Path(anidir),
                                                       "TriangularWave",
                                                       self.frequency,
                                                       80000, 1000)


class SineWave(CommonWave):
    def __init__(self, frequency=440.0, loudness=0.5, samplerate=44100, decay=0.996):
        """
        :param frequency: 1秒内有多少周期
        :param samplerate: 1秒内有多少数据点
        :param loudness: 响度
        :param decay: 声音是否自然消散
        """
        super(SineWave, self).__init__(round(samplerate / frequency), loudness, decay)
        self.frequency = frequency
        self.samplerate = samplerate

    def pluck(self):
        self.buffer.clear()
        # PureTone的频率是用数学公式严谨定义的，不用担心是不是不对
        for i in range(self.init_capacity):
            self.buffer.append(self.loudness * math.sin(2 * math.pi * self.frequency * i / self.samplerate))

    def gen_wave_ani(self, anidir: str):
        super(SineWave, self).gen_wave_ani_super(Path(anidir),
                                                 "SineWave",
                                                 self.frequency,
                                                 80000, 1000)

    def gen_wave_frame(self):
        """
        测试PIL绘制图像功能，这里将不会继续开发
        :return:
        """
        w, h = 128, 64
        zplt = ZPlot(w, h, "1")
        cut = math.ceil(self.init_capacity / w)
        y = list(self.buffer)[::cut]
        zplt.plot(y)
        return zplt.get_frame()


class Instrument(WaveBuffer):
    def __init__(self, init_capacity,
                 decay=0.996,
                 factor: float = 1.0,
                 samplerate=44100,
                 beta=0.5):
        """
        :param init_capacity: 一个周期的采样点数量
        :param decay: 消散率
        :param factor:
        :param samplerate: 采样率，1秒内有多少数据点
        :param beta:
        """
        super(Instrument, self).__init__(init_capacity)
        self.buffer = None
        self.init_capacity = init_capacity
        self.decay = decay
        self.factor = factor
        self.samplerate = samplerate
        self.beta = beta

    def pluck(self):
        self.buffer = deque()
        for i in range(self.init_capacity):
            self.buffer.append(random.random() - 0.5)

    def tic(self):
        """
        执行一次Karplus Strong算法
        :return:
        """
        a1 = self.buffer.popleft()
        a2 = self.buffer[0]
        if self.factor > random.random():
            # new_a = self.DECAY * 0.5 * (a1 + a2)
            new_a = self.decay * (self.beta * a1 + (1 - self.beta) * a2)
        else:
            new_a = -self.decay * (self.beta * a1 + (1 - self.beta) * a2)
        self.buffer.append(new_a)

    def play(self, length=80000):
        """
        一次性播放，播放时无法做其他事情
        :param length:
        :return:
        """
        audio = StdAudio()
        for i in range(length):
            audio.play_float(self.sample())
            self.tic()


class GuitarString(Instrument):
    def __init__(self, frequency=440.0):
        SR = 44100
        super(GuitarString, self).__init__(round(SR / frequency),
                                           0.996, 1, SR, 0.1)
        self.frequency = frequency
        self.capacity = round(self.samplerate / frequency)
        super(GuitarString, self).pluck()

    def gen_wave_ani(self, anidir: str):
        super(GuitarString, self).gen_wave_ani_super(Path(anidir),
                                                     "GuitarString",
                                                     self.frequency,
                                                     80000, 1000)

    def play(self, length=80000):
        self.pluck()
        super(GuitarString, self).play(length)


class Harp(Instrument):
    def __init__(self, frequency=440.0):
        SR = 44100
        super(Harp, self).__init__(round(SR / frequency),
                                   decay=0.998,
                                   factor=0,
                                   samplerate=SR,
                                   beta=0.5
                                   )
        self.frequency = frequency
        super(Harp, self).pluck()


class Drum(Instrument):
    def __init__(self, frequency=440.0):
        SR = 44100
        super(Drum, self).__init__(round(SR / frequency),
                                   decay=1,
                                   factor=0.5,
                                   samplerate=SR,
                                   beta=0.5
                                   )
        self.frequency = frequency
        super(Drum, self).pluck()


def main():
    f = 220
    gs = GuitarString(f)
    gs.gen_wave_ani(f'../../output/ani/gs_{f}.mp4')
    # while True:
    #     gs.play()

    # h = Harp(f)
    # h.play()

    # d = Drum(110)
    # d.play()


if __name__ == "__main__":
    main()
